#include "ANTLRv4ParserBase.hpp"
#include "cpp-gen/ANTLRv4Parser.h"
#include <functional>
#include <optional>

using namespace antlr4;

// TODO Find the reason why there is such strange identifier
static const StringList ReservedLocalnames = {
    "type",
    "aliasType",
    "isArrayList",
    "<missing ARG_ACTION_WORD>"};

static auto isReservedLocalName(const std::string &key)
{
    for (const auto &str : ReservedLocalnames)
        if (key == str)
            return true;
    return false;
}

ANTLRv4ParserBase::ANTLRv4ParserBase(antlr4::TokenStream *input) : Parser(input)
{
}

ANTLRv4ParserBase::~ANTLRv4ParserBase()
{
}

static const auto formatRuleName = [](const std::string &ruleName)
{
    auto copy = ruleName;
    if (copy.size())
        copy[0] = toupper(copy.at(0));
    return copy;
};

std::string ANTLRv4ParserBase::parserRuleNameToContextName(const std::string &ruleName)
{
    auto parserRuleName = ruleName;
    if (parserRuleName.size())
        parserRuleName[0] = toupper(parserRuleName.at(0));
    return std::string(parserRuleName + "Context");
}

bool ANTLRv4ParserBase::ruleIsParserRule(ParserRuleContext *context)
{
    return dynamic_cast<ANTLRv4Parser::ParserRuleSpecContext *>(context);
}

const RuleCollection &ANTLRv4ParserBase::ruleCollection() const
{
    return rulesVisibleToVisitor_;
}

const RuleCollection &ANTLRv4ParserBase::invisibleRules() const
{
    return rulesInvisibleToVisitor_;
}

static auto removeSuffixInElement(std::string &element)
{
    if (element.size() && element.back() == '_')
        element.resize(element.size() - 1);
    return element;
};

// NOTE type && aliasType are used as reserved local names
static auto localsInRuleSpec(ANTLRv4Parser::ParserRuleSpecContext *ctx)
{
    if (ctx->localsSpec() && ctx->localsSpec()->argActionParameters() && ctx->localsSpec()->argActionParameters()->argActionParameter().size())
    {
        StringList localsList;
        for (const auto localParam : ctx->localsSpec()->argActionParameters()->argActionParameter())
            if (localParam->name && (!isReservedLocalName(localParam->name->getText())))
                localsList.push_back(localParam->name->getText());
        if (localsList.size())
            return std::make_optional(localsList);
    }
    return std::optional<StringList>();
}

void ANTLRv4ParserBase::findParserRule(ParserRuleContext *context)
{
    auto ruleContext = static_cast<ANTLRv4Parser::ParserRuleSpecContext *>(context);
    bool ruleAltIHasPound = false; // Indicating if the rule 's alternative is followed with "#LabelName"
    auto localsInRule = localsInRuleSpec(ruleContext);
    if (localsInRule.has_value())
    {
        std::cout << ruleContext->name->getText() << ": ";
        for (const auto &str : localsInRule.value())
            std::cout << str << " ";
        std::cout << std::endl;
        std::cout.flush();
    }

    const auto insertRuleElement = [&](BaseRuleInfor &details, ANTLRv4Parser::LabeledAltContext *alt)
    {
        if (alt->alternative() && alt->alternative()->elements())
            for (const auto &element : alt->alternative()->elements()->element())
            {
                if (element->labeledElement() && element->labeledElement()->label)
                    details.ruleElementNames.push_back(element->labeledElement()->label->getText());
            }
        details.locals = (localsInRule.has_value()) ? localsInRule.value() : StringList();
        rulesVisibleToVisitor_.insert(std::make_pair(details.ruleName, std::make_shared<BaseRuleInfor>(details)));
    };

    for (const auto &alt : static_cast<ANTLRv4Parser::ParserRuleSpecContext *>(context)->ruleBlock()->ruleAltList()->labeledAlt())
    {
        if (alt->hasPOUND)
            ruleAltIHasPound = true;
        if (auto altId = alt->id())
        {
            BaseRuleInfor ruleNameAndElements;
            ruleNameAndElements.ruleName = formatRuleName(altId->getText());
            ruleNameAndElements.ruleContextName = formatRuleName(parserRuleNameToContextName(altId->getText()));
            insertRuleElement(ruleNameAndElements, alt);
        }
    }

    if (ruleAltIHasPound)
    {
        BaseRuleInfor ruleNameAndElements;
        ruleNameAndElements.ruleName = formatRuleName(ruleContext->name->getText());
        ruleNameAndElements.ruleContextName =
            parserRuleNameToContextName(ruleContext->name->getText());
        rulesInvisibleToVisitor_.insert(std::make_pair(ruleNameAndElements.ruleName, std::make_shared<BaseRuleInfor>(ruleNameAndElements)));
    }
    else
    {
        BaseRuleInfor ruleNameAndElements;
        ruleNameAndElements.ruleName = formatRuleName(ruleContext->name->getText());
        ruleNameAndElements.ruleContextName =
            parserRuleNameToContextName(ruleContext->name->getText());
        ruleNameAndElements.locals = localsInRule.has_value() ? localsInRule.value() : StringList();
        insertRuleElement(ruleNameAndElements, ruleContext->ruleBlock()->ruleAltList()->labeledAlt().at(0));
    }
}

void ANTLRv4ParserBase::reset()
{
    rulesVisibleToVisitor_.clear();
    rulesInvisibleToVisitor_.clear();
}
